package org.finesys.chat.core.base.embedding.factory;

import lombok.extern.slf4j.Slf4j;
import org.finesys.ai.entity.AigcKnowledge;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
public class KnowledgeStoreFactory {

    private final Map<String, AigcKnowledge> knowledgeMap = new ConcurrentHashMap<>();

    public KnowledgeStoreFactory() {
    }

    /**
     * 初始话方法
     */
    public void init(List<AigcKnowledge> knowledgeList) {
        // 停止之前的嵌入存储服务
        if (!knowledgeMap.isEmpty()) {
            knowledgeMap.clear();
        }
        //重新初始化嵌入存储服务
        knowledgeList.forEach(this::addOrUpdate);
    }

    /**
     * 添加或更新嵌入存储服务
     *
     * @param knowledge 嵌入存储配置对象
     */
    public void addOrUpdate(AigcKnowledge knowledge) {
        //重新初始化嵌入存储服务
        String id = knowledge.getId();
        knowledgeMap.put(id, knowledge);
    }

    public void unregister(AigcKnowledge knowledge) {
        String id = knowledge.getId();
        if (knowledgeMap.containsKey(id)) {
            knowledgeMap.remove(id);
        }
    }

    /**
     * 根据给定的id获取对应的嵌入存储。
     *
     * @param id 嵌入存储的唯一标识符
     * @return 与给定id对应的嵌入存储对象，如果id不存在，则返回null
     */
    public AigcKnowledge get(String id) {
        return knowledgeMap.get(id);
    }

    public boolean contains(String id) {
        return knowledgeMap.containsKey(id);
    }
}
